import os
import sys
import json
from util import ArgumentParser, ConfigLoader, LOG
from model import OpenAIModel
from translator import PDFTranslator
from flask import Flask, request, flash, redirect, Response, render_template, send_file
from werkzeug.utils import secure_filename
from flask_cors import CORS

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

app = Flask(__name__)
CORS(app)

global translator


def init_app():
    argument_parser = ArgumentParser()
    args = argument_parser.parse_arguments()
    config_loader = ConfigLoader(args.config)
    config = config_loader.load_config()

    model_name = args.openai_model if args.openai_model else config['OpenAIModel']['model']
    api_key = args.openai_api_key if args.openai_api_key else config['OpenAIModel']['api_key']

    model = OpenAIModel(model=model_name, api_key=api_key)
    # 实例化 PDFTranslator 类，并调用 translate_pdf() 方法
    global translator
    translator = PDFTranslator(model)


app.config['UPLOAD_FOLDER'] = '.\\tests\\uploadFiles'
app.config['TRANSLATOR_FOLDER'] = 'tests\\translatedFiles'
app.config['ALLOWED_EXTENSIONS'] = {'pdf'}


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']


@app.route('/')
def home():
    return render_template('translator.html')


@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        flash('没有文件部分')
        return redirect(request.url)
    file = request.files['file']
    if file.filename == '':
        flash('没有选择文件')
        return redirect(request.url)
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        file.save(os.path.join(app.config['UPLOAD_FOLDER'], filename))
        return Response("文件上传成功", status=200)
    else:
        return Response("文件格式不支持", status=415)


format_dict = {
    "markdown": ".md",
    "pdf": ".pdf",
}


@app.route('/translate', methods=['POST'])
def translate():
    file = request.files['file']
    target_language = request.form.get('target_language')
    target_format = request.form.get('target_format')

    if not target_language:
        return 'Missing target language.', 400

    if not target_format:
        return 'Missing target format.', 400

    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        upload_file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        # 先保存到本地
        file.save(upload_file_path)

        try:
            translatedFileName = filename.split(".")[0] + "_" + target_language + format_dict[target_format]
            output_file_path = os.path.join(app.config['TRANSLATOR_FOLDER'], translatedFileName)
            translator.translate_pdf(upload_file_path, target_format, target_language, output_file_path)
            data = {
                "SUCCESS": True,
                "message": "翻译完成！",
                "translatedFileName": translatedFileName
            }
            return Response(json.dumps(data), status=200)
        except Exception as e:
            LOG.error(f"An error occurred during translation: {e}")
            return Response("{\"message\": \"翻译过程异常！\"}", status=500)

    else:
        return Response("文件格式不支持", status=415)


@app.route('/download', methods=['GET'])
def download_file():
    filename = request.args.get('filename')
    return send_file(os.path.join(app.config['TRANSLATOR_FOLDER'], filename), as_attachment=True)


if __name__ == "__main__":
    init_app()
    app.run()
